import json
from typing import TypedDict, List, Dict, Tuple, Generator
import networkx as nx
from clang import cindex
from MelodieFuncFlow import MelodieGenerator
import logging

logger = logging.getLogger(__name__)


class Position(TypedDict):
    startLine: int
    startCol: int
    endLine: int
    endCol: int


class CFGNode(TypedDict):
    id: int
    position: Position
    terminatorKind: str


class Edge(TypedDict):
    source: int
    target: int
    order: int


class FunctionCFG(TypedDict):
    functionName: str
    functionPosition: Position
    nodes: List[CFGNode]
    edges: List[Edge]


class ClangCFG(nx.DiGraph):
    def statements(self, node_id: int) -> List[cindex.Cursor]:
        return self.nodes[node_id].get("statements", [])


CONTROL_FLOW_CURSORS = [
    cindex.CursorKind.IF_STMT,
    cindex.CursorKind.SWITCH_STMT,
    cindex.CursorKind.FOR_STMT,
    cindex.CursorKind.WHILE_STMT,
    cindex.CursorKind.DO_STMT,
    cindex.CursorKind.CASE_STMT,
    cindex.CursorKind.DEFAULT_STMT,
    cindex.CursorKind.RETURN_STMT,
    cindex.CursorKind.CONTINUE_STMT,
    cindex.CursorKind.COMPOUND_STMT,
    cindex.CursorKind.BREAK_STMT,
]


def _walker(node: cindex.Cursor) -> Generator[cindex.Cursor, None, None]:
    for child in node.get_children():
        if child.kind not in CONTROL_FLOW_CURSORS:  # Not a control flow node
            yield child
        else:
            yield from _walker(child)


def recursive_walk_non_control_flow_nodes(
    cu: cindex.Cursor,
) -> MelodieGenerator[cindex.Cursor]:
    return MelodieGenerator(_walker(cu))


class CFGManager:
    def __init__(self, cfgs_file: str, _end_line=99999):
        with open(cfgs_file) as f:
            self._cfgs_data: List[FunctionCFG] = json.load(f)
        self._topologies = [
            self._build_nx_graph(cfg_data) for cfg_data in self._cfgs_data
        ]
        self._end_line = _end_line
        self._nodes_range: List[Tuple[int, int, Position]] = []
        for cfg_index, cfg_data in enumerate(self._cfgs_data):
            for node in cfg_data["nodes"]:
                self._nodes_range.append((cfg_index, node["id"], node["position"]))

    def _locate_line_col_to_node(self, line_col: Tuple[int, int]) -> Tuple[int, int]:
        """
        Given a line and column, locate the node in the CFG and the cfg index in the `_cfgs_data.nodes`.
        Returns a tuple (cfg_index, node_id) if found, otherwise (None, None).
        """
        line, col = line_col

        for cfg_idx, node_id, position in self._nodes_range:
            if (
                position["startLine"] <= line <= position["endLine"]
                and position["startCol"] <= col <= position["endCol"]
            ):
                return (cfg_idx, node_id)

        return (None, None)

    def _find_function_cfg_index(self, start_line: int, start_col: int) -> int:
        """
        Find the CFG data for the function containing the given line and column.
        """
        for cfg_idx, cfg_data in enumerate(self._cfgs_data):
            function_position = cfg_data["functionPosition"]
            if (
                function_position["startLine"] == start_line
                and function_position["startCol"] == start_col
            ):
                return cfg_idx
        return -1

    def _locate_cursor_to_cfg_node(
        self, func_cfg_index: int, cu: cindex.Cursor
    ) -> Tuple[int]:
        """
        Locate a statement to a function's CFG node

        If not found, return -1
        """

        line, col = cu.location.line, cu.location.column

        for node_data in self._cfgs_data[func_cfg_index]["nodes"]:
            node_id = node_data["id"]
            position = node_data["position"]
            if (
                position["startLine"] <= line <= position["endLine"]
                and position["startCol"] <= col <= position["endCol"]
            ):
                return node_id

        return -1

    def _build_nx_graph(self, cfg_data: FunctionCFG) -> ClangCFG:
        """
        将 JSON 格式的 CFG 转换为 ClangCFG。
        每个节点包含 terminatorKind 属性。
        每条边包含 source, target 和 order 属性。
        """
        G = ClangCFG()

        # 添加节点
        for node in cfg_data["nodes"]:
            G.add_node(
                node["id"],
                terminator_kind=node["terminatorKind"],  # 添加 terminatorKind 属性
                position=node["position"],  # 可选：添加位置信息
            )

        # 添加边
        for edge in cfg_data["edges"]:
            G.add_edge(
                edge["source"],
                edge["target"],
                order=edge["order"],  # 边属性：order，意思是顺序。
            )

        return G

    def construct_cfg_with_stmts(
        self, cfg_id: int, func_decl_cu: cindex.Cursor
    ) -> ClangCFG:
        """
        Construct CFG with statements.
        """
        new_cfg = self._topologies[cfg_id].copy()
        for cu in recursive_walk_non_control_flow_nodes(func_decl_cu):
            node_id = self._locate_cursor_to_cfg_node(cfg_id, cu)
            if node_id == -1:
                logging.warning(
                    f"Cannot locate cursor to CFG node: {cu.kind} {cu.location.line}, {cu.location.column}"
                )
                continue
            print(cu.location.line, cu.location.column, cu.kind, node_id)
            if new_cfg.nodes[node_id].get("statements") is None:
                new_cfg.nodes[node_id]["statements"] = []
            new_cfg.nodes[node_id]["statements"].append(cu)
        return new_cfg

    @staticmethod
    def format_label(cfg: ClangCFG):
        cfg_repr = cfg.copy()
        for node_id, node in cfg_repr.nodes(data=True):
            # if node.get("statements") is not None:
            statements_str = "\n".join(
                [
                    " ".join([t.spelling for t in s.get_tokens()])
                    + f"  line {s.location.line}:{s.location.column}"
                    for s in node.get("statements", [])
                ]
            )
            node["label"] = (
                f"#{node_id} {node.get('terminator_kind', '')}\n{statements_str}"
            )
        return cfg_repr
